import os
import torch
from tqdm import tqdm
import numpy as np
from loadDataset import MakeDataSet
from modelSource import Multi_featureV2 as multi_person_feature


def train(rootPath: str):
    A_noTargetData, A_targetData = MakeDataSet.LoadTarAndNoTarData(os.path.join(rootPath, "MyDataSet/SubjectA"), True)
    B_noTargetData, B_targetData = MakeDataSet.LoadTarAndNoTarData(os.path.join(rootPath, "MyDataSet/SubjectB"), True)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model = multi_person_feature.multi_person_feature(1, 2)
    model = model.to(device)

    # 损失函数
    loss_fn = torch.nn.CrossEntropyLoss()
    # 优化器
    learning_rate = 0.01
    optim = torch.optim.SGD(model.parameters(), learning_rate, momentum=0.9)

    if not os.path.exists(os.path.join(rootPath, "model")):
        os.mkdir(os.path.join(rootPath, "model"))
    modelSavePath = os.path.join(rootPath, "model")
    # 训练的轮数
    epoch = 100

    for i in tqdm(range(epoch)):
        # print("当前Epoch：{}".format(i))
        model.train()
        for j in range(np.size(A_targetData, 0) - 10):
            A_electrodes = A_targetData[j, 61:64, :, :, :]
            B_electrodes = B_targetData[j, 61:64, :, :, :]
            for k in range(np.size(A_electrodes, 1)):
                for m in range(np.size(A_electrodes, 2)):
                    A_data = A_electrodes[:, k, m, :]
                    B_data = B_electrodes[:, k, m, :]
                    data = np.concatenate((A_data, B_data))
                    if np.all(data[0] == 0):
                        continue
                    data = torch.from_numpy(data).reshape(1, 1, 240, 6).transpose(2, 3)
                    label = torch.tensor([1], dtype=torch.long).reshape(1)
                    data = data.to(device)
                    label = label.to(device)
                    # print(data)
                    output = model(data)
                    # print(output.argmax(1))
                    # print(label)
                    loss = loss_fn(output, label)  # label需要Long类型
                    loss.backward()
                    optim.step()

        for j in range(np.size(A_noTargetData, 0) - 10):
            A_electrodes = A_noTargetData[j, 61:64, :, :, :]
            B_electrodes = B_noTargetData[j, 61:64, :, :, :]
            for k in range(np.size(A_electrodes, 1)):
                for m in range(np.size(A_electrodes, 2)):
                    A_data = A_electrodes[:, k, m, :]
                    B_data = B_electrodes[:, k, m, :]
                    data = np.concatenate((A_data, B_data))
                    if np.all(data[0] == 0):
                        continue
                    data = torch.from_numpy(data).reshape(1, 1, 240, 6).transpose(2, 3)
                    label = torch.tensor([0], dtype=torch.long).reshape(1)
                    data = data.to(device)
                    label = label.to(device)
                    # print(data)
                    output = model(data)
                    # print(output.argmax(1))
                    # print(label)
                    loss = loss_fn(output, label)  # label需要Long类型
                    loss.backward()
                    optim.step()
        total_test_loss = 0
        total_accuracy = 0
        model.eval()
        totalCount = 0
        with torch.no_grad():
            for e in range(np.size(A_targetData, 0) - 10, np.size(A_targetData, 0)):
                for j in range(np.size(A_targetData, 0) - 10):
                    A_electrodes = A_targetData[e, 61:64, :, :, :]
                    B_electrodes = B_targetData[e, 61:64, :, :, :]
                    for k in range(np.size(A_electrodes, 1)):
                        for m in range(np.size(A_electrodes, 2)):
                            A_data = A_electrodes[:, k, m, :]
                            B_data = B_electrodes[:, k, m, :]
                            data = np.concatenate((A_data, B_data))
                            if np.all(data[0] == 0):
                                continue
                            data = torch.from_numpy(data).reshape(1, 1, 240, 6).transpose(2, 3)
                            label = torch.tensor([1], dtype=torch.long).reshape(1)
                            data = data.to(device)
                            label = label.to(device)
                            output = model(data)
                            loss = loss_fn(output, label)
                            # total_train_step+=1
                            total_test_loss += loss
                            # print(loss)
                            accuracy = (output.argmax(1) == label).sum().item()
                            total_accuracy += accuracy
                            totalCount += 1
            for e in range(np.size(A_noTargetData, 0) - 10, np.size(A_noTargetData, 0)):
                for j in range(np.size(A_noTargetData, 0) - 10):
                    A_electrodes = A_noTargetData[e, 61:64, :, :, :]
                    B_electrodes = B_noTargetData[e, 61:64, :, :, :]
                    for k in range(np.size(A_electrodes, 1)):
                        for m in range(np.size(A_electrodes, 2)):
                            A_data = A_electrodes[:, k, m, :]
                            B_data = B_electrodes[:, k, m, :]
                            data = np.concatenate((A_data, B_data))
                            if np.all(data[0] == 0):
                                continue
                            data = torch.from_numpy(data).reshape(1, 1, 240, 6).transpose(2, 3)
                            label = torch.tensor([0], dtype=torch.long).reshape(1)
                            data = data.to(device)
                            label = label.to(device)
                            output = model(data)
                            loss = loss_fn(output, label)
                            # total_train_step+=1
                            total_test_loss += loss
                            # print(loss)
                            accuracy = (output.argmax(1) == label).sum().item()
                            total_accuracy += accuracy
                            totalCount += 1
        print("Epoch：{}，整体测试集上的Loss:{}".format(i, total_test_loss))
        print("Epoch：{}，整体测试集上的Accuray:{}".format(i, total_accuracy / totalCount))
        torch.save(model, os.path.join(modelSavePath, "test_{}_{:5f}.pth".format(i, total_accuracy / totalCount)))
        # total_test_step += 1
